from PIL import Image
import torch
from torch.utils.data import Dataset
import cv2
import os
import pandas as pd
import numpy as np
from trans import get_train_transform

class AffectDataSet2(Dataset):
    def __init__(self, data_path, phase, affcls, transform=None):
        self.phase = phase
        self.transform = transform
        self.data_path = data_path
        self.cls = affcls

        if affcls == 7:
            if phase == 'train':
                df = pd.read_csv(os.path.join(self.data_path, '7cls_train.txt'), sep=' ', header=None,names=['name','label'])
            else:
                df = pd.read_csv(os.path.join(self.data_path, '7cls_val.txt'), sep=' ', header=None,names=['name','label'])
        else:
            if phase == 'train':
                df = pd.read_csv(os.path.join(self.data_path, '8cls_train.txt'), sep=' ', header=None,names=['name','label'])
            else:
                df = pd.read_csv(os.path.join(self.data_path, '8cls_val.txt'), sep=' ', header=None,names=['name','label'])

        file_names = df["name"]
        self.data = df
        self.label = self.data.loc[:, 'label'].values

        _, self.sample_counts = np.unique(self.label, return_counts=True)
        print(f' distribution of {phase} samples: {self.sample_counts}')

        self.file_paths = []
        if phase == 'train':
            for f in file_names:
                path = os.path.join(self.data_path, 'trainnew', f)
                self.file_paths.append(path)
        else:
            for f in file_names:
                path = os.path.join(self.data_path, 'validnew', f)
                self.file_paths.append(path)

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        path = self.file_paths[idx]
        image = cv2.imread(path)
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # OpenCV读取的是BGR，需要转换为RGB
        label = self.label[idx]

        if self.transform is not None:
            augmented = self.transform(image=image)
            image = augmented['image']  # 提取tensor
            
            # 转换为单通道图像
            # 由于ToGray已经使所有通道值相同，只需保留第一个通道即可
            # if image.shape[0] == 3:  # 如果是3通道
            #     image = image[0:1, :, :]  # 只保留第一个通道，维度变为[1, H, W]

        return image, label

    def get_labels(self):
        return self.label

if __name__ == "__main__":
    train_transform = get_train_transform()
    test_dataset = AffectDataSet2(data_path = "/Users/tunm/Downloads/AffectNet", affcls=7, phase = 'train', transform=train_transform)
    print(len(test_dataset))
    print(test_dataset[0][0].shape)